#include "sd_memory.h"
#include "sd_blist.h"
#include "sd_log.h"
#include "sd_thread.h"
#include "sd_hashtable.h"

typedef struct allocRecord
{
	void* addr;
	size_t size;
	uint64 id;
	uint32 tid;
	uint64 thread_block_id;
}allocRecord;

static sd_MemoryAllocator g_DebugAllocator;
static sd_MemoryAllocator g_SystemAllocator;
static bool g_IsMemoryIntialized = false;
static sd_SpinLock g_AllocatorLock;
static void* sd_DebugAllocate(size_t size);
static void  sd_DebugDeallocate(void* p);
static uint64  g_AllocIDCounter = 0;
static sd_Hashtable* g_pThreadAllocIDCounters = NULL;

sd_Ret sd_MemoryInit(void)
{
	g_SystemAllocator.allocate = malloc;
	g_SystemAllocator.deallocate = free;
	
	g_DebugAllocator.allocate = sd_DebugAllocate;
	g_DebugAllocator.deallocate = sd_DebugDeallocate;
	sd_BlistCreate(&g_DebugAllocator.allocRecords,sizeof(allocRecord), 10240, &g_SystemAllocator);
	sd_HashtableCreate(&g_pThreadAllocIDCounters, 
		sd_HashFunc_uint32, 
		sd_HashKeyCompare_uint32, 
		7,
		5, 
		sizeof(uint32), 
		sizeof(uint64), 
		&g_SystemAllocator);

	sd_SpinLock_Init(&g_AllocatorLock);

	g_IsMemoryIntialized = true;
	return SD_RET_OK;
}

sd_Ret sd_MemoryExit(void)
{
	sd_BlistFree(&g_DebugAllocator.allocRecords);
	
	sd_HashtableFree(g_pThreadAllocIDCounters);

	sd_SpinLock_Destroy(&g_AllocatorLock);
	g_IsMemoryIntialized = false;
	return SD_RET_OK;
}

sd_Ret sd_MemoryIsInitialized(void)
{
	return (sd_Ret)g_IsMemoryIntialized;
}

sd_Ret sd_MemoryCheck(void)
{
	size_t leakedMemsize = 0;
	sd_BlistIterator iter;
	
	sd_SpinLock_Lock(&g_AllocatorLock);
	
	sd_BlistBegin(&g_DebugAllocator.allocRecords, &iter);
	while(sd_BlistNext(&iter) != SD_RET_FALSE)
	{
		allocRecord rc;
		sd_BlistGet(&iter, &rc);
		sd_log_err("memory block %lld in thread %d 's block %lld is not freed!\n",rc.id, rc.tid, rc.thread_block_id);
		leakedMemsize += rc.size;
	}
	sd_log_err("Totally %d bytes memory is not freed!\n",leakedMemsize);
	SD_ASSERT(leakedMemsize == 0);
	
	sd_SpinLock_Unlock(&g_AllocatorLock);
	return SD_RET_OK;
}

sd_MemoryAllocator* sd_GetDefaultMemoryAllocator()
{
#ifdef _DEBUG_MEMORY
	return &g_DebugAllocator;
#else
	return &g_SystemAllocator;
#endif
}

void* sd_malloc(size_t size)
{
	return sd_GetDefaultMemoryAllocator()->allocate(size);
}

void sd_free(void* p)
{
	sd_GetDefaultMemoryAllocator()->deallocate(p);
}

static void* sd_DebugAllocate(size_t size)
{
	uint32 tid = sd_getCurrentThreadID();
	allocRecord rc;
	uint64 tid_counter;
	void *p;
	sd_SpinLock_Lock(&g_AllocatorLock);

	if (sd_HashtableLookup(g_pThreadAllocIDCounters, &tid, &tid_counter) != SD_RET_OK)
	{
		tid_counter=0;
		sd_HashtableInsert(g_pThreadAllocIDCounters, &tid, &tid_counter);
	}
	
	p = g_SystemAllocator.allocate(size);
	rc.addr = p;
	rc.size = size;
	rc.id = g_AllocIDCounter++;
	rc.tid = tid;
	rc.thread_block_id = tid_counter++;

	sd_BlistInsert(&g_DebugAllocator.allocRecords, &rc);
	sd_HashtableUpdate(g_pThreadAllocIDCounters, &tid, &tid_counter);
	
	sd_SpinLock_Unlock(&g_AllocatorLock);
	return p;
}

static void sd_DebugDeallocate(void* p)
{
	sd_BlistIterator iter;
	bool bBlockFound=false;
	sd_SpinLock_Lock(&g_AllocatorLock);

	sd_BlistBegin(&g_DebugAllocator.allocRecords, &iter);
	while(sd_BlistNext(&iter) != SD_RET_FALSE)
	{
		allocRecord rc;
		sd_BlistGet(&iter, &rc);
		if (rc.addr == p)
		{
			g_SystemAllocator.deallocate(p);
			sd_BlistRemove(&iter);
			bBlockFound = true;
			break;
		}
	}
	if (!bBlockFound)
	{
		g_SystemAllocator.deallocate(p);
		sd_log_err("Not found this memory block %x in all records!", (int)p); 
		SD_ASSERT(false); /* Not found this memory block in all records */
	}
	sd_SpinLock_Unlock(&g_AllocatorLock);
}

